import torch
import torch.nn as nn
import numpy as np
from tqdm import tqdm
from models.scheduler import get_schedule_jump

class DiffusionModel(nn.Module):
    """
    RePaint: 使用重采样的去噪扩散概率模型进行图像修复。
    该模型在扩散过程中保持未掩码区域不变，并使用跳跃式重采样来改善生成质量。
    """
    def __init__(
        self,
        model,
        beta_start=1e-4,
        beta_end=0.02,
        timesteps=1000,
        device="cuda"
    ):
        super().__init__()
        self.model = model  # UNet模型，用于预测噪声
        self.timesteps = timesteps  # 扩散步数
        self.device = device
        
        # 定义beta时间表（噪声水平逐步增加）
        self.betas = torch.linspace(beta_start, beta_end, timesteps, device=device)
        
        # 定义alphas及其累积乘积（噪声相关参数）
        self.alphas = 1. - self.betas
        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
        self.alphas_cumprod_prev = torch.cat([torch.tensor([1.0], device=device), self.alphas_cumprod[:-1]])
        
        # 计算扩散q(x_t | x_{t-1})和其他辅助变量
        self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)
        self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1. - self.alphas_cumprod)
        self.log_one_minus_alphas_cumprod = torch.log(1. - self.alphas_cumprod)
        self.sqrt_recip_alphas_cumprod = torch.sqrt(1. / self.alphas_cumprod)
        self.sqrt_recipm1_alphas_cumprod = torch.sqrt(1. / self.alphas_cumprod - 1)
        
        # 计算后验q(x_{t-1} | x_t, x_0)
        self.posterior_variance = self.betas * (1. - self.alphas_cumprod_prev) / (1. - self.alphas_cumprod)
        self.posterior_log_variance_clipped = torch.log(torch.max(self.posterior_variance, torch.tensor(1e-20, device=device)))
        self.posterior_mean_coef1 = self.betas * torch.sqrt(self.alphas_cumprod_prev) / (1. - self.alphas_cumprod)
        self.posterior_mean_coef2 = (1. - self.alphas_cumprod_prev) * torch.sqrt(self.alphas) / (1. - self.alphas_cumprod)
    
    def q_sample(self, x_0, t, noise=None):
        """
        对图像进行前向扩散过程，将噪声添加到x_0得到x_t
        
        Args:
            x_0: 原始图像
            t: 时间步
            noise: 可选的预定义噪声
            
        Returns:
            x_t: t时刻的噪声图像
        """
        if noise is None:
            noise = torch.randn_like(x_0)
            
        sqrt_alphas_cumprod_t = self.sqrt_alphas_cumprod[t].reshape(-1, 1, 1, 1)
        sqrt_one_minus_alphas_cumprod_t = self.sqrt_one_minus_alphas_cumprod[t].reshape(-1, 1, 1, 1)
        
        return sqrt_alphas_cumprod_t * x_0 + sqrt_one_minus_alphas_cumprod_t * noise
    
    def p_losses(self, x_0, t, mask=None, noise=None):
        """
        训练损失计算
        
        Args:
            x_0: 原始图像
            t: 时间步
            mask: 修复区域的掩码，1表示需要修复的区域
            noise: 可选的预定义噪声
            
        Returns:
            loss: 训练损失
        """
        if noise is None:
            noise = torch.randn_like(x_0)
            
        # 获取噪声图像
        x_t = self.q_sample(x_0, t, noise)
        
        # 预测噪声（epsilon预测模式）
        predicted_noise = self.model(x_t, t)
        
        # 计算损失
        if mask is not None:
            # 只在掩码区域计算损失
            loss = torch.mean((noise - predicted_noise)**2 * mask)
        else:
            loss = torch.mean((noise - predicted_noise)**2)
            
        return loss

    def forward(self, x_0, mask=None):
        """
        训练前向传播
        
        Args:
            x_0: 原始图像
            mask: 修复区域的掩码，1表示需要修复的区域
            
        Returns:
            loss: 训练损失
        """
        # 随机选择时间步
        t = torch.randint(0, self.timesteps, (x_0.shape[0],), device=self.device).long()
        return self.p_losses(x_0, t, mask)
    
    @torch.no_grad()
    def p_sample_ddpm(self, x_t, t, t_index):
        """
        标准DDPM单步去噪采样：从x_t采样得到x_{t-1}
        
        Args:
            x_t: 当前噪声图像
            t: 当前时间步
            t_index: 时间步索引
            
        Returns:
            x_{t-1}: 去噪一步后的图像
        """
        betas_t = self.betas[t].reshape(-1, 1, 1, 1)
        sqrt_one_minus_alphas_cumprod_t = self.sqrt_one_minus_alphas_cumprod[t].reshape(-1, 1, 1, 1)
        sqrt_recip_alphas_t = self.sqrt_recip_alphas_cumprod[t].reshape(-1, 1, 1, 1)
        
        # 模型预测噪声
        predicted_noise = self.model(x_t, t)
        
        # 计算均值 - 预测x_0
        pred_x0 = sqrt_recip_alphas_t * x_t - sqrt_recip_alphas_t * sqrt_one_minus_alphas_cumprod_t * predicted_noise
        pred_x0 = torch.clamp(pred_x0, -1.0, 1.0)
        
        # 计算后验方差和随机噪声
        posterior_variance_t = self.posterior_variance[t].reshape(-1, 1, 1, 1)
        noise = torch.randn_like(x_t) if t_index > 0 else torch.zeros_like(x_t)
        
        # 计算x_{t-1}，即下一步去噪后的结果
        x_t_minus_1 = (
            self.posterior_mean_coef1[t].reshape(-1, 1, 1, 1) * pred_x0 +
            self.posterior_mean_coef2[t].reshape(-1, 1, 1, 1) * x_t +
            torch.sqrt(posterior_variance_t) * noise
        )
        
        return x_t_minus_1
    
    @torch.no_grad()
    def p_sample_repaint(self, x_t, t, mask, original_image):
        """
        RePaint采样：结合已知区域和未知区域
        
        Args:
            x_t: 当前噪声图像
            t: 当前时间步（标量）
            mask: 修复区域的掩码，1表示需要修复的区域
            original_image: 原始图像（未掩码部分）
            
        Returns:
            合并后的图像
        """
        # 对未知区域进行标准DDPM采样
        x_unknown = self.p_sample_ddpm(x_t, torch.tensor([t], device=self.device).long(), t)
        
        # 对已知区域，从原始图像采样相应噪声水平的版本
        if t > 0:
            noise = torch.randn_like(original_image)
            x_known = self.q_sample(original_image, torch.tensor([t-1], device=self.device).long(), noise)
        else:
            x_known = original_image
        
        # 合并已知和未知区域
        x_t_minus_1 = mask * x_unknown + (1 - mask) * x_known
        
        return x_t_minus_1
    
    @torch.no_grad()
    def p_sample_loop_repaint(self, shape, mask, original_image, 
                              num_resample_U=10, 
                              jump_length=10, 
                              jump_n_sample=10,
                              jump_start_resampling_step_T_fraction=0.0,
                              use_repaint_schedule_fig9_approx=True):
        """
        RePaint采样循环，使用跳跃式重采样策略
        
        Args:
            shape: 输出图像形状
            mask: 修复区域的掩码，1表示需要修复的区域
            original_image: 原始图像（未掩码部分）
            num_resample_U: 每一步的重采样次数（论文中的U）
            jump_length: 跳跃长度（论文中的j）
            jump_n_sample: 每个跳跃点的重采样次数（论文中的r）
            jump_start_resampling_step_T_fraction: 开始重采样的时间步比例
            use_repaint_schedule_fig9_approx: 是否使用论文图9的调度近似
            
        Returns:
            生成的修复图像
        """
        device = self.device
        batch_size = shape[0]
        
        # 从纯噪声开始
        img = torch.randn(shape, device=device)
        
        # 获取时间步调度
        if use_repaint_schedule_fig9_approx:
            # 使用论文中的跳跃调度
            start_resampling = int(self.timesteps * (1 - jump_start_resampling_step_T_fraction))
            time_steps = get_schedule_jump(
                t_T=self.timesteps-1,
                n_sample=num_resample_U,
                jump_length=jump_length,
                jump_n_sample=jump_n_sample,
                start_resampling=start_resampling
            )
        else:
            # 简单的线性调度用于测试
            time_steps = list(reversed(range(0, self.timesteps)))
        
        # 跟踪当前时间步
        prev_t = self.timesteps - 1
        
        # 逐步去噪
        for i, t_now in enumerate(tqdm(time_steps, desc="RePaint采样")):
            t_now = int(t_now)
            
            if t_now < 0:  # 结束标记
                break
                
            # 处理时间步的变化
            if t_now > prev_t:  # 向前跳跃（重采样）
                # 从x_{t-1}重新加噪到x_t
                noise = torch.randn_like(img)
                img = self.q_sample_from_to(img, prev_t, t_now, noise)
            elif t_now < prev_t:  # 正常的去噪步骤
                # 从x_t去噪到x_{t_now}
                for t in reversed(range(t_now, prev_t)):
                    img = self.p_sample_repaint(img, t, mask, original_image)
            
            prev_t = t_now
        
        # 最终的合并
        final_known = original_image
        return mask * img + (1 - mask) * final_known
    
    def q_sample_from_to(self, x_from, t_from, t_to, noise=None):
        """
        从时间步t_from的图像采样到时间步t_to的图像
        用于重采样过程中的前向跳跃
        """
        if noise is None:
            noise = torch.randn_like(x_from)
        
        # 计算从t_from到t_to的转换
        alpha_from = self.alphas_cumprod[t_from] if t_from >= 0 else 1.0
        alpha_to = self.alphas_cumprod[t_to]
        
        # 计算转换系数
        sqrt_alpha_to_from = torch.sqrt(alpha_to / alpha_from)
        sqrt_one_minus_alpha_to_from = torch.sqrt((1 - alpha_to) - (1 - alpha_from) * (alpha_to / alpha_from))
        
        # 应用转换
        return sqrt_alpha_to_from * x_from + sqrt_one_minus_alpha_to_from * noise
    
    @torch.no_grad()
    def sample(self, shape, mask, original_image):
        """
        生成修复后的图像样本（为了兼容性保留的接口）
        """
        return self.p_sample_loop_repaint(shape, mask, original_image)
    
    def inpaint(self, image, mask, 
                num_resample_U=10,
                jump_length=10, 
                jump_n_sample=10,
                jump_start_resampling_step_T_fraction=0.0,
                use_repaint_schedule_fig9_approx=True):
        """
        执行图像修复（RePaint方法）
        
        Args:
            image: 带有缺失区域的原始图像，值范围[-1, 1]
            mask: 修复区域的掩码，1表示需要修复的区域，0表示保持原样
            num_resample_U: 每一步的重采样次数
            jump_length: 跳跃长度
            jump_n_sample: 每个跳跃点的重采样次数
            jump_start_resampling_step_T_fraction: 开始重采样的时间步比例
            use_repaint_schedule_fig9_approx: 是否使用论文的调度策略
            
        Returns:
            修复后的图像
        """
        shape = image.shape
        
        # 使用RePaint方法生成修复后的图像
        inpainted_image = self.p_sample_loop_repaint(
            shape, mask, image,
            num_resample_U=num_resample_U,
            jump_length=jump_length,
            jump_n_sample=jump_n_sample,
            jump_start_resampling_step_T_fraction=jump_start_resampling_step_T_fraction,
            use_repaint_schedule_fig9_approx=use_repaint_schedule_fig9_approx
        )
        
        return inpainted_image